import cv2
import numpy as np
from flask import current_app
from numba import njit
from numba.typed import List
import torch
from app.models.steel_defect_models import SegModel
from config import Config  # 从配置文件导入配置
import json
import os
import matplotlib.pyplot as plt
from fpdf import FPDF
import logging
from pathlib import Path
from tqdm import tqdm

@njit
def _mask2pixels(flattened_mask: np.ndarray, class_id: int) -> List[int]:
    fill_val = class_id
    pixels = List()
    idx = 0
    start_idx = 0
    count = 0
    for i in flattened_mask:
        if i == fill_val:
            if count == 0:
                start_idx = idx
            count += 1
        else:
            if count > 0:
                pixels.append(start_idx + 1)
                pixels.append(count)
                count = 0
        idx += 1
    if count > 0:
        pixels.append(start_idx + 1)
        pixels.append(count)
    return pixels

def mask2pixels(mask: np.ndarray, class_id: int) -> List[int]:
    return _mask2pixels(mask.flatten("F"), class_id)

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def load_models():
    config = Config()  # 实例化配置类
    models = []
    checkpoint_paths = tqdm(config.CHECKPOINT_PATHS, desc="Loading models", unit="model")
    for ckpt in checkpoint_paths:
        try:
            model = SegModel.load_from_checkpoint(
                checkpoint_path=ckpt,
                map_location=config.DEVICE,
                arch="FPN",
                encoder_name="resnet34",
                in_channels=1,
                out_classes=5,
                t_max=config.EPOCHS * config.BATCH_SIZE,
                mean=torch.tensor(config.MEAN, dtype=torch.float),
                std=torch.tensor(config.STD, dtype=torch.float)
            ).to(config.DEVICE).eval()
            models.append(model)
        except Exception as e:
            logging.error(f"Failed to load model from checkpoint {ckpt}: {str(e)}", exc_info=True)
    return models

# 修改 process_image 函数，直接使用全局变量
def process_image(image_path):
    global preloaded_models
    from app.model_loader import preloaded_models
    if preloaded_models is None:
        logging.error("Models are not available for processing the image.")
        return None
    try:
        # 读取图片，支持 JPEG、PNG、BMP 等常见格式
        img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
        if img is None:
            raise ValueError(f"Failed to load image: {image_path}")
        img_tensor = torch.tensor(img, dtype=torch.float, device="cpu").unsqueeze(0).unsqueeze(0)

        # 多模型预测并投票
        with torch.no_grad():
            masks = []
            for model in preloaded_models:
                try:
                    mask = model(img_tensor).contiguous().softmax(dim=1).argmax(dim=1).squeeze()
                    masks.append(mask)
                except Exception as e:
                    logging.error(f"Error in model prediction for {image_path}: {str(e)}")
            if not masks:
                raise Exception(f"No valid predictions for {image_path}")
            masks = torch.stack(masks)
            pred_mask = torch.mode(masks, dim=0).values.cpu().numpy()

        # 后处理
        mask_cc = np.zeros_like(pred_mask)
        for class_id in range(1, 5):
            num_component, component = cv2.connectedComponents((pred_mask == class_id).astype(np.uint8))
            for c in range(1, num_component):
                p = (component == c)
                if p.sum() > 2000:  # 过滤小区域
                    mask_cc[p] = class_id
        return mask_cc
    except Exception as e:
        logging.error(f"Error processing image {image_path}: {str(e)}", exc_info=True)
        return None

def generate_report(input_urls_json, mask_urls_json, passed_percent, sum_qty):
    input_urls = json.loads(input_urls_json)
    mask_urls = json.loads(mask_urls_json)
    failed_percent = 1 - passed_percent

    # 绘制饼状图
    labels = ['Passed', 'Failed']
    sizes = [passed_percent, failed_percent]
    plt.pie(sizes, labels=labels, autopct='%1.1f%%')
    plt.axis('equal')
    plt.title('Pass/Fail Ratio')
    pie_chart_path = f'static/pie_chart.png'
    plt.savefig(pie_chart_path)
    plt.close()

    # 绘制柱状图
    plt.bar(labels, sizes)
    plt.xlabel('Result')
    plt.ylabel('Percentage')
    plt.title('Pass/Fail Ratio')
    bar_chart_path = f'static/bar_chart.png'
    plt.savefig(bar_chart_path)
    plt.close()

    report = {
        '输入图片链接': input_urls,
        '掩码图片链接': mask_urls,
        '合格率': passed_percent
    }
    report_json = json.dumps(report, ensure_ascii=False)

    # 生成 PDF 报告
    pdf = FPDF()
    pdf.add_page()

    # 获取当前脚本所在目录
    current_script_dir = Path(__file__).parent

    # 构建字体文件的相对路径
    font_path = current_script_dir.parent / 'fonts' / 'simhei.ttf'

    try:
        pdf.add_font('SimHei', '', str(font_path), uni=True)
        pdf.set_font('SimHei', size=12)
    except Exception as e:
        logging.error(f"字体文件 {font_path} 未找到或添加失败: {e}，使用默认字体。")
        pdf.set_font("Arial", size=12)

    pdf.cell(200, 10, txt="检测报告", ln=True, align='C')

    # 添加饼状图及文字介绍
    pdf.cell(200, 10, txt="饼状图展示了检测结果的合格率和不合格率分布", ln=True, align='C')
    pdf.ln(5)  # 留出一定空间

    # 计算饼状图居中的 x 坐标
    page_width = pdf.w
    chart_width = 100
    pie_x = (page_width - chart_width) / 2
    pdf.image(pie_chart_path, x=pie_x, y=pdf.get_y(), w=chart_width)
    pdf.ln(100)  # 增加饼状图和柱状图之间的间距

    # 添加柱状图及文字介绍
    pdf.cell(200, 10, txt="柱状图直观对比了检测结果的合格率和不合格率", ln=True, align='C')
    pdf.ln(5)  # 留出一定空间

    # 计算柱状图居中的 x 坐标
    bar_x = (page_width - chart_width) / 2
    pdf.image(bar_chart_path, x=bar_x, y=pdf.get_y(), w=chart_width)
    pdf.ln(100)  # 留出足够的空间给图表

    # 添加检测结果图片
    for index, (input_url, mask_url) in enumerate(zip(input_urls, mask_urls)):
        app_root_path = current_app.root_path
        input_path = os.path.join(app_root_path, '..', input_url[1:])
        mask_path = os.path.join(app_root_path, '..', mask_url[1:])

        pdf.cell(200, 10, txt=f"图片 {index + 1} 检测结果", ln=True, align='C')
        pdf.cell(100, 10, txt=f"输入图片 {index + 1}", ln=False, align='C')
        pdf.cell(100, 10, txt=f"掩码图片 {index + 1}", ln=True, align='C')

        try:
            pdf.image(input_path, x=10, y=pdf.get_y(), w=80)
            pdf.image(mask_path, x=100, y=pdf.get_y(), w=80)
            pdf.ln(80)  # 留出足够的空间给图片
        except Exception as e:
            logging.error(f"Error adding images to PDF for result at index {index}: {e}")

    # 获取当前脚本所在目录
    current_script_dir = Path(__file__).parent
    # 获取项目根目录
    project_root = current_script_dir.parent
    project_root = project_root.parent

    # 构建 PDF 文件的绝对路径
    pdf_path = project_root / 'static' / 'prediction_report.pdf'

    try:
        pdf.output(str(pdf_path))
        logging.info(f"PDF report generated successfully: {pdf_path}")
    except Exception as e:
        logging.error(f"Failed to generate PDF report: {str(e)}", exc_info=True)
        return None, None, None, None

    return report_json, str(pdf_path), passed_percent, report_json















